package uk.ac.cam.bt288.Pico2;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.Cipher;

public class Keys {
	private static int pubKeyBytes = 162;
	private static int privKeyBytes = 636;
	private static SecureRandom random = new SecureRandom();

	public static String nextSessionId()
	{
		return new BigInteger(256, random).toString(32);
	}

	public static String formatKey (Key k) {
		StringBuffer format = new StringBuffer();
		format.append("Algorithm: " + k.getAlgorithm() + ", Format: " + k.getFormat() +
				", Bytes: " + k.getEncoded().length + "\n");

		format.append("Key: " + new BigInteger(1, k.getEncoded()).toString(16) + "\n");

		return format.toString();
	}

	public static KeyPair GenerateKeyPair () {
		KeyPairGenerator gen = null;
		KeyPair keyPair = null;
		try {
			gen = KeyPairGenerator.getInstance("RSA"); //generate RSA key pair
			gen.initialize(1024); //key size
		}
		catch (Exception e) {
			e.printStackTrace();
		}
		try {
			keyPair = gen.generateKeyPair();
		}
		catch (Exception e) {
			e.printStackTrace();
		}

		return keyPair;
	}

	public static void SaveKey (KeyPair k, String app) throws IOException {
		//write public key
		FileOutputStream fout = new FileOutputStream(app +".pub"); //don't append
		byte[] pubKey = k.getPublic().getEncoded();
		fout.write(pubKey);
		fout.close();

		//write private key
		FileOutputStream privFout = new FileOutputStream(app +".priv");
		byte[] privKey = k.getPrivate().getEncoded();
		privFout.write(privKey);
		privFout.close();


	}

	public static KeyPair readKeyPair (String app) throws IOException, NoSuchAlgorithmException, InvalidKeySpecException {
		FileInputStream pubFin = new FileInputStream (app + ".pub");
		FileInputStream priFin = new FileInputStream (app + ".priv");

		byte[] encodedPub = new byte[pubKeyBytes];
		pubFin.read(encodedPub);
		byte[] encodedPriv = new byte[privKeyBytes];
		priFin.read(encodedPriv);
		KeyFactory rsaKeyFac = KeyFactory.getInstance("RSA");
		PublicKey pubKey = rsaKeyFac.generatePublic(new X509EncodedKeySpec(encodedPub));
		PrivateKey priKey = rsaKeyFac.generatePrivate(new PKCS8EncodedKeySpec(encodedPriv));
		return new KeyPair(pubKey, priKey);
	}

	/*
	public static KeyPair[] readKeyPair (String app) throws IOException, NoSuchAlgorithmException, InvalidKeySpecException {
		File pubKeyFile = new File(app + ".pub");
		FileInputStream pubFin = new FileInputStream (pubKeyFile);
		FileInputStream privFin = new FileInputStream (app + ".priv");

		int num = (int) (pubKeyFile.length() / (8 * pubKeyBytes));
		KeyPair[] rtn = new KeyPair[num];

		System.err.println(num);

		for (int i=0; i<num; i++) {
			byte[] encodedPub = new byte[pubKeyBytes];
			pubFin.read(encodedPub, i*pubKeyBytes, pubKeyBytes);

			System.err.println(new BigInteger(1, encodedPub).toString(16));

			KeyFactory dsaKeyFac = KeyFactory.getInstance("DSA");
			PublicKey pubKey = dsaKeyFac.generatePublic(new X509EncodedKeySpec(encodedPub));

			byte[] encodedPriv = new byte[privKeyBytes];
			privFin.read(encodedPriv, i*privKeyBytes, privKeyBytes);

			System.err.println(new BigInteger(1, encodedPriv).toString(16));

			PrivateKey privateKey = dsaKeyFac.generatePrivate(new PKCS8EncodedKeySpec(encodedPriv));

			KeyPair k = new KeyPair(pubKey, privateKey);
			rtn[i] = k;
		}

		pubFin.close();
		privFin.close();

		return rtn;

	}
	 */

	public static byte[] encrypt(String s, Key k) throws Exception {
		Cipher cipher = Cipher.getInstance("RSA/ECB/NoPadding");
		cipher.init(Cipher.ENCRYPT_MODE, k);
		byte[] cipherText = cipher.doFinal(s.getBytes());
		return cipherText;
	}

	public static String decrypt(byte[] data, Key k) throws Exception {
		Cipher cipher = Cipher.getInstance("RSA/ECB/NoPadding");
		cipher.init(Cipher.DECRYPT_MODE, k);
		byte[] plainText = cipher.doFinal(data);
		return new String(plainText);
	}

	public static void main (String[] args) throws Exception {
		KeyPair k = GenerateKeyPair();
		try {
			SaveKey(k, "test1");
			//SaveKey(GenerateKeyPair(), "test1");
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}

		System.out.println("Read:");

		KeyPair keys = readKeyPair("test1");
		System.out.println(formatKey(keys.getPublic()));
		System.out.println(formatKey(keys.getPrivate()));

		String plainText = "1234567890";
		byte[] data = encrypt(plainText, k.getPublic());
		System.err.println(data.length);
		System.out.println("Encrypted: " + new String(data));
		System.out.println("Plain    : " + decrypt(data, k.getPrivate()));
		System.out.println();

		byte[] data2 = encrypt(plainText, k.getPrivate());
		System.err.println(data2.length);
		System.out.println("Encrypted: " + new String(data2));
		System.out.println("Decrypted: " + decrypt(data2, k.getPublic()));
	}
}
